#include <limits>
#include <cassert>
#include <cmath>
#include "../common/common.hpp"
#include <TYImageProc.h>


static TY_CAMERA_INTRINSIC m_colorIntrinsic;
DepthViewer depthViewer0("diffDepth");
DepthViewer depthViewer1("filteredDepth");
// DepthViewer depthViewer2("diffDepth");


static int cvpf2typf(int cvpf)
{
    switch(cvpf){
        case CV_8U: return TY_PIXEL_FORMAT_MONO;
        case CV_8UC3: return TY_PIXEL_FORMAT_RGB;
        case CV_16UC1: return TY_PIXEL_FORMAT_DEPTH16;
        default: return TY_PIXEL_FORMAT_UNDEFINED;
    }
}

static void mat2TY_IMAGE_DATA(int comp, const cv::Mat& mat, TY_IMAGE_DATA& data)
{
    data.status = 0;
    data.componentID = comp;
    data.size = mat.total() * mat.elemSize();
    data.buffer = mat.data;
    data.width = mat.cols;
    data.height = mat.rows;
    data.pixelFormat = cvpf2typf(mat.type());
}


struct CallbackData {
    int             index;
    TY_DEV_HANDLE   hDevice;
    PCLPointCloudViewer* pcviewer;
    TY_CAMERA_INTRINSIC* intri;

    std::vector<cv::Mat> depths;
    size_t          depthNum;

    bool saveOneFramePoint3d;
    int  fileIndex;
};

cv::Mat depthToWorld(float* intr, const cv::Mat &depth)
{
    cv::Mat world(depth.rows, depth.cols, CV_32FC3);
    float cx = intr[2];
    float cy = intr[5];
    float inv_fx = 1.0f / intr[0];
    float inv_fy = 1.0f / intr[4];
    for (int r = 0; r < depth.rows; r++)
    {
        uint16_t* pSrc = (uint16_t*)depth.data + r * depth.cols;
        cv::Vec3f* pDst = (cv::Vec3f*)world.data + r * depth.cols;
        for (int c = 0; c < depth.cols; c++)
        {
            uint16_t z = pSrc[c];
            if(z == 0){
                  pDst[c][0] = NAN;
                  pDst[c][1] = NAN;
                  pDst[c][2] = NAN;
            } else {
                  pDst[c][0] = (c - cx) * z * inv_fx;
                  pDst[c][1] = (r - cy) * z * inv_fy;
                  pDst[c][2] = z;
            }
        }
    }
    return world;
}


void frameHandler(TY_FRAME_DATA* frame, void* userdata)
{
    CallbackData* pData = (CallbackData*) userdata;
    LOGD("=== Get frame %d", ++pData->index);

    cv::Mat depth, color, p3d;
    parseFrame(*frame, &depth, 0, 0, &color);
    if(pData->saveOneFramePoint3d){
        LOGE("No point3d now!!!");
        /*
        char file[32];
        sprintf(file, "points-%d.xyz", pData->fileIndex++);
        writePointCloud((cv::Point3f*)p3d.data, p3d.total(), file, PC_FILE_FORMAT_XYZ);
        pData->saveOneFramePoint3d = false;
        imshow("depth", depth * 32);
        */
    }

    if(!depth.empty()){
        pData->depths.push_back(depth.clone());
        if(pData->depths.size() >= pData->depthNum){
          // filter
          std::vector<TY_IMAGE_DATA> tyDepth(pData->depthNum);
          for(size_t i = 0; i < pData->depthNum; i++){
            mat2TY_IMAGE_DATA(TY_COMPONENT_DEPTH_CAM, pData->depths[i], tyDepth[i]);
          }
          TY_IMAGE_DATA tyFilteredDepth;
          cv::Mat filteredDepth(depth.size(), depth.type());
          mat2TY_IMAGE_DATA(TY_COMPONENT_DEPTH_CAM, filteredDepth, tyFilteredDepth);

          struct DepthEnhenceParameters param = DepthEnhenceParameters_Initializer;
          // param.sigma_s = 10;
          // param.sigma_r = 20;
          // param.outlier_rate = 0.95;
          // param.outlier_win_sz = 20;
          TYDepthEnhenceFilter(&tyDepth[0], pData->depthNum, NULL, &tyFilteredDepth, &param);
          pData->depths.clear();

          cv::Mat orgDepth = filteredDepth.clone();
          struct DepthSpeckleFilterParameters sfparam = DepthSpeckleFilterParameters_Initializer;
          TYDepthSpeckleFilter (&tyFilteredDepth, &sfparam);
          cv::Mat diffDepth = orgDepth - filteredDepth;
          depthViewer0.show(diffDepth);
          depthViewer1.show(filteredDepth);

          cv::Mat newP3d = depthToWorld(pData->intri->data, filteredDepth);
          pData->pcviewer->show(newP3d, "newPoint3D");

          // cv::Mat diff = depth - filteredDepth;
          // depthViewer2.show("diffDepth", diff);
        }
    }
    if(!color.empty()){
        imshow("Color", color);
    }

    LOGD("=== Re-enqueue buffer(%p, %d)", frame->userBuffer, frame->bufferSize);
    ASSERT_OK( TYEnqueueBuffer(pData->hDevice, frame->userBuffer, frame->bufferSize) );
}

int main(int argc, char* argv[])
{
    std::string ID, IP;
    TY_INTERFACE_HANDLE hIface = NULL;
    TY_DEV_HANDLE hDevice = NULL;
    TY_CAMERA_INTRINSIC intri;
    int depthNum = 3;

    for(int i = 1; i < argc; i++){
        if(strcmp(argv[i], "-id") == 0){
            ID = argv[++i];
        } else if(strcmp(argv[i], "-ip") == 0) {
            IP = argv[++i];
        } else if(strcmp(argv[i], "-d") == 0) {
            depthNum = atoi(argv[++i]);
        }else if(strcmp(argv[i], "-h") == 0){
            LOGI("Usage: SimpleView_Callback [-h] [-id <ID>]");
            return 0;
        }
    }

    LOGD("=== Init lib");
    ASSERT_OK( TYInitLib() );
    TY_VERSION_INFO ver;
    ASSERT_OK( TYLibVersion(&ver) );
    LOGD("     - lib version: %d.%d.%d", ver.major, ver.minor, ver.patch);

    std::vector<TY_DEVICE_BASE_INFO> selected;
    ASSERT_OK( selectDevice(TY_INTERFACE_ALL, ID, IP, 1, selected) );
    ASSERT(selected.size() > 0);
    TY_DEVICE_BASE_INFO& selectedDev = selected[0];

    ASSERT_OK( TYOpenInterface(selectedDev.iface.id, &hIface) );
    ASSERT_OK( TYOpenDevice(hIface, selectedDev.id, &hDevice) );

    LOGD("=== Configure components, open depth cam");
    int32_t componentIDs = TY_COMPONENT_DEPTH_CAM;
    ASSERT_OK( TYEnableComponents(hDevice, componentIDs) );

    //try to enable depth map
    LOGD("Configure components, open depth cam");
    if (componentIDs & TY_COMPONENT_DEPTH_CAM) {
        int32_t image_mode;
        ASSERT_OK(get_default_image_mode(hDevice, TY_COMPONENT_DEPTH_CAM, image_mode));
        LOGD("Select Depth Image Mode: %dx%d", TYImageWidth(image_mode), TYImageHeight(image_mode));
        ASSERT_OK(TYSetEnum(hDevice, TY_COMPONENT_DEPTH_CAM, TY_ENUM_IMAGE_MODE, image_mode));
        ASSERT_OK(TYEnableComponents(hDevice, TY_COMPONENT_DEPTH_CAM));
    }

    LOGD("=== Get depth intrinsic");
    ASSERT_OK( TYGetStruct(hDevice, TY_COMPONENT_DEPTH_CAM, TY_STRUCT_CAM_INTRINSIC, &intri, sizeof(intri)) );

    LOGD("=== Prepare image buffer");
    uint32_t frameSize;
    ASSERT_OK( TYGetFrameBufferSize(hDevice, &frameSize) );
    LOGD("     - Get size of framebuffer, %d", frameSize);
    ASSERT( frameSize >= 640*480*2 );

    LOGD("     - Allocate & enqueue buffers");
    char* frameBuffer[2];
    frameBuffer[0] = new char[frameSize];
    frameBuffer[1] = new char[frameSize];
    LOGD("     - Enqueue buffer (%p, %d)", frameBuffer[0], frameSize);
    ASSERT_OK( TYEnqueueBuffer(hDevice, frameBuffer[0], frameSize) );
    LOGD("     - Enqueue buffer (%p, %d)", frameBuffer[1], frameSize);
    ASSERT_OK( TYEnqueueBuffer(hDevice, frameBuffer[1], frameSize) );

    LOGD("=== Disable trigger mode");
    TY_TRIGGER_PARAM trigger;
    trigger.mode = TY_TRIGGER_MODE_OFF;
    ASSERT_OK(TYSetStruct(hDevice, TY_COMPONENT_DEVICE, TY_STRUCT_TRIGGER_PARAM, &trigger, sizeof(trigger)));

    LOGD("=== Start capture");
    ASSERT_OK( TYStartCapture(hDevice) );

    PCLPointCloudViewer pcviewer;
    CallbackData cb_data;
    cb_data.index = 0;
    cb_data.hDevice = hDevice;
    cb_data.pcviewer = &pcviewer;
    cb_data.depthNum = depthNum;
    cb_data.saveOneFramePoint3d = false;
    cb_data.fileIndex = 0;
    cb_data.intri = &intri;

    LOGD("=== While loop to fetch frame");
    TY_FRAME_DATA frame;
    bool exit_main = false;
    while(!exit_main){
        int key = cv::waitKey(100);
        switch(key & 0xff){
            case 0xff:
                break;
            case 'q':
                exit_main = true;
                break;
            case 's':
                cb_data.saveOneFramePoint3d = true;
                break;
            default:
                LOGD("Pressed key %d", key);
        }

        int err = TYFetchFrame(hDevice, &frame, -1);
        if( err != TY_STATUS_OK ){
            LOGD("... Drop one frame");
            continue;
        }

        frameHandler(&frame, &cb_data);
    }

    ASSERT_OK( TYStopCapture(hDevice) );
    ASSERT_OK( TYCloseDevice(hDevice) );
    ASSERT_OK( TYCloseInterface(hIface) );
    ASSERT_OK( TYDeinitLib() );
    delete frameBuffer[0];
    delete frameBuffer[1];

    LOGD("=== Main done!");
    return 0;
}
